import subprocess
import sys
import pkg_resources
import platform

# 定义镜像源列表
MIRROR_SOURCES = [
    'https://pypi.tuna.tsinghua.edu.cn/simple',  # 清华
    'https://mirrors.aliyun.com/pypi/simple',    # 阿里
    'https://pypi.mirrors.ustc.edu.cn/simple',   # 中科大
    'https://mirrors.cloud.tencent.com/pypi/simple'  # 腾讯
]

def select_mirror():
    """选择镜像源"""
    print("\n可用的镜像源:")
    for i, mirror in enumerate(MIRROR_SOURCES, 1):
        print(f"{i}. {mirror}")
    print("0. 使用默认源(PyPI)")
    
    while True:
        try:
            choice = int(input("\n请选择镜像源 (0-4): "))
            if 0 <= choice <= len(MIRROR_SOURCES):
                if choice == 0:
                    return None
                return MIRROR_SOURCES[choice-1]
        except ValueError:
            pass
        print("无效的选择，请重试")

def check_python_version():
    """检查Python版本"""
    print("\n检查Python版本...")
    python_version = sys.version.split()[0]
    print(f"当前Python版本: {python_version}")
    if sys.version_info < (3, 7):
        print("警告: 推荐使用Python 3.7或更高版本")
        return False
    return True

def get_required_packages():
    """获取所需的包列表"""
    return {
        'tensorflow': '2.10.0',
        'opencv-python': '4.5.0',
        'numpy': '1.19.2',
        'matplotlib': '3.3.2',
        'pillow': '8.0.0',
        'scikit-learn': '0.24.0',
        'tqdm': '4.50.0'
    }

def check_gpu():
    """检查GPU支持"""
    print("\n检查GPU支持...")
    try:
        import tensorflow as tf
        gpus = tf.config.list_physical_devices('GPU')
        if gpus:
            print(f"找到 {len(gpus)} 个GPU设备:")
            for gpu in gpus:
                print(f"- {gpu.device_type}: {gpu.name}")
            return True
        else:
            print("未找到GPU设备，将使用CPU模式")
            return False
    except Exception as e:
        print(f"检查GPU时出错: {str(e)}")
        return False

def install_package(package_name, version=None, mirror=None):
    """安装指定的包"""
    if version:
        package_spec = f"{package_name}=={version}"
    else:
        package_spec = package_name
    
    print(f"\n正在安装 {package_spec}...")
    try:
        cmd = [sys.executable, "-m", "pip", "install", package_spec]
        if mirror:
            cmd.extend(["-i", mirror, "--trusted-host", mirror.split('/')[2]])
        subprocess.check_call(cmd)
        return True
    except subprocess.CalledProcessError as e:
        print(f"安装 {package_spec} 失败: {str(e)}")
        return False

def check_and_install_packages():
    """检查并安装所需的包"""
    required_packages = get_required_packages()
    installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
    
    print("\n检查所需的Python包...")
    packages_to_install = {}
    
    for package, version in required_packages.items():
        if package not in installed_packages:
            print(f"未安装: {package}")
            packages_to_install[package] = version
        else:
            current_version = installed_packages[package]
            print(f"已安装: {package} (版本 {current_version})")
            if current_version != version:
                print(f"版本不匹配，需要: {version}")
                packages_to_install[package] = version
    
    if packages_to_install:
        print("\n需要安装以下包:")
        for package, version in packages_to_install.items():
            print(f"- {package} (版本 {version})")
        
        choice = input("\n是否现在安装这��包？(y/n): ")
        if choice.lower() == 'y':
            mirror = select_mirror()
            for package, version in packages_to_install.items():
                install_package(package, version, mirror)
        else:
            print("已取消安装")
    else:
        print("\n所有必需的包都已正确安装！")

def check_system_info():
    """检查系统信息"""
    print("\n系统信息:")
    print(f"操作系统: {platform.system()} {platform.version()}")
    print(f"处理器: {platform.processor()}")
    print(f"Python位置: {sys.executable}")
    
    # 安全地获取pip版本信息
    try:
        pip_cmd = [sys.executable, '-m', 'pip', '--version']
        pip_version = subprocess.check_output(pip_cmd, stderr=subprocess.STDOUT, text=True)
        print(f"pip信息: {pip_version.strip()}")
    except subprocess.CalledProcessError as e:
        print("警告: 无法获取pip信息，可能需要安装pip")
        choice = input("是否要安装pip? (y/n): ")
        if choice.lower() == 'y':
            try:
                # 下载get-pip.py
                print("正在下载pip安装脚本...")
                subprocess.check_call([
                    sys.executable, 
                    "-c", 
                    "import urllib.request; urllib.request.urlretrieve('https://bootstrap.pypa.io/get-pip.py', 'get-pip.py')"
                ])
                # 安装pip
                print("正在安装pip...")
                subprocess.check_call([sys.executable, "get-pip.py"])
                print("pip安装成功！")
            except Exception as e:
                print(f"安装pip失败: {str(e)}")
                print("请手动安装pip后重试")
                sys.exit(1)
    except Exception as e:
        print(f"警告: 检查pip时出错: {str(e)}")

def main():
    """主函数"""
    print("=== 环境检查工具 ===")
    
    # 检查系统信息
    check_system_info()
    
    # 检查Python版本
    if not check_python_version():
        choice = input("\n是否继续？(y/n): ")
        if choice.lower() != 'y':
            return
    
    # 检查并安装包
    check_and_install_packages()
    
    # 检查GPU支持
    check_gpu()
    
    print("\n环境检查完成！")

if __name__ == "__main__":
    main() 